1139C - Edgy Trees - CodeForces Solution


dfs and similar dsu graphs math trees *1500

Please click on ads to support us..

Python Code:

import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def get_root(s):
    v = []
    while not s == root[s]:
        v.append(s)
        s = root[s]
    for i in v:
        root[i] = s
    return s

def unite(s, t):
    rs, rt = get_root(s), get_root(t)
    if not rs ^ rt:
        return
    if rank[s] == rank[t]:
        rank[rs] += 1
    if rank[s] >= rank[t]:
        root[rt] = rs
        size[rs] += size[rt]
    else:
        root[rs] = rt
        size[rt] += size[rs]
    return

def same(s, t):
    return True if get_root(s) == get_root(t) else False

def get_size(s):
    return size[get_root(s)]

n, k = map(int, input().split())
mod = pow(10, 9) + 7
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
for _ in range(n - 1):
    u, v, x = map(int, input().split())
    if not x:
        unite(u, v)
ng = 0
for i in range(1, n + 1):
    if i ^ get_root(i):
        continue
    ng += pow(get_size(i), k, mod)
    ng %= mod
ans = pow(n, k, mod) - ng
ans %= mod
print(ans)

C++ Code:

//......................................// Work hard in silence, let the success make the noise //..............................//
//.....................................................//THINK TWICE .. CODE ONCE//.............................................//
#include<bits/stdc++.h> 
#define improve ios_base::sync_with_stdio(false); cin.tie(0), cout.tie(0);  
#define ll long long 
#define ld long double
#define dv vector<ll> 
#define dvp vector<pair<ll,ll>> 
#define dit vector<ll>::iterator 
#define dvv vector<vector<ll>>
#define dvb vector<bool>
#define dm map<ll,ll>
#define ds set<ll>
#define loop(n) for (ll i = 0; i < n; i++)
#define loopj(n) for (ll j = 0; j < n; j++)
#define loopit(v) for (auto it = v.begin(); it != v.end(); ++it)
#define input(n,v) loop(n) { ll in; cin >> in; v.push_back(in);} 
#define inputvp(n,v) loop(n) {ll x,y; cin>>x>>y; v.push_back(make_pair(x,y));}
#define prefix_sum(n,v,s) loop(n) { s[i] = v[i]; if (i)s[i] += s[i - 1]; }
#define suffix_sum(n,v,s) reverse(all(v)); prefix_sum(n, v, s) reverse(all(s))
#define all(v) v.begin(), v.end()
#define bs(v,n) binary_search(all(v), n)
#define lb(v,x) lower_bound(all(v), x) - v.begin()
#define ub(v,x) upper_bound(all(v), x) - v.begin()
#define fix(n, res) fixed <<  setprecision(n) << (ld)res
#define posmod(n,m) ((n % m)+ m) % m
#define chk(ret) for (auto& v : ret)cout << v <<" "; cout<<"\n"
#define chkvp(ret) for (auto& v : ret)cout << v.first<<" "<< v.second << "\n";
#define chkmp(m)  for (auto const &pair: m) cout << "{" << pair.first << " : " << pair.second << "}\n"
#define chkmpv(m)     for(auto const &m: mp){cout << m.first << " : ";for (auto const &s : m.se)  cout << s << " ";cout << "\n"; }
#define sz(s) s.size() 
#define pb push_back
#define  EPS = 1e-10
#define fi first
#define se second  
#define pi 3.14159265358979323 
using namespace std;
const ll MOD = 1e9 + 7, N = 1e6 + 5, MAX = 500;

dv par(N, 0);
dv s(N, 1);
ll find_parent(ll u)
{
    if (par[u] == u)return u;
    else return par[u] = find_parent(par[u]);
}
bool is_con(ll u, ll v)
{
    u = find_parent(u);
    v = find_parent(v);
    return (u == v);
}
void con(ll u, ll v)
{
    if (is_con(u, v))return;
    u = find_parent(u);
    v = find_parent(v); 
    if (s[u] < s[v])
    { 
        par[u] = v;
        s[v] += s[u];
    }
    else
    { 
        par[v] = u;
        s[u] += s[v];
    }
}
ll int binpow(ll int a, ll int b, ll int m) {
    a %= m;
    ll int res = 1;
    while (b > 0)
    {
        if (b & 1)
            res = res * a % m;
        a = a * a % m;
        b >>= 1;
    }
    return res;
}
int main()
{
    improve;
    loop(N)par[i] = i;
    ll n = 0, k = 0; cin >> n >> k;
    loop(n - 1)
    {
        ll u = 0, v = 0, c = 0; cin >> u >> v >> c; 
        u--, v--;
        if (!c)con(u, v);
    }  
    set<ll>st;
    loop(n) 
        st.insert(find_parent(par[i])); 
    ll cnt = 0;
    for (auto i : st)
    {
        cnt = (cnt+ binpow(s[i], k, MOD))%MOD ; 
       
    }
    cout << (binpow(n, k, MOD) - cnt + MOD) % MOD;

}
	 		     	 		     		 	  	 	  		


Comments

Submit
0 Comments
More Questions

376A - Lever
1305A - Kuroni and the Gifts
1609A - Divide and Multiply
149B - Martian Clock
205A - Little Elephant and Rozdil
1609B - William the Vigilant
978B - File Name
1426B - Symmetric Matrix
732B - Cormen --- The Best Friend Of a Man
1369A - FashionabLee
1474B - Different Divisors
1632B - Roof Construction
388A - Fox and Box Accumulation
451A - Game With Sticks
768A - Oath of the Night's Watch
156C - Cipher
545D - Queue
459B - Pashmak and Flowers
1538A - Stone Game
1454C - Sequence Transformation
165B - Burning Midnight Oil
17A - Noldbach problem
1350A - Orac and Factors
1373A - Donut Shops
26A - Almost Prime
1656E - Equal Tree Sums
1656B - Subtract Operation
1656A - Good Pairs
1367A - Short Substrings
87A - Trains